import torch
from PIL import Image
from transformers import AutoModelForCausalLM, AutoTokenizer
import os
device = "cuda"
model_path = "/rainbow/mnt/ai_pro/sd/swift/glm_4v_9b"


model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
    trust_remote_code=True
).to(device).eval()

    
prompt_txt = f"""
如图是一个购物小票，请从图中按顺序识别以下 4 个项目：
1.商户名称（注意商户名称在图中没有字段标记，可能是商户的品牌名或商铺名）；
2.订单编号（注意订单编号在图中可能被标记为“交易号”、“交易序号”、“流水号”、“收据”、“订单号”、“账单号”、“账单”或“No”字段，请根据此字段顺序提取一个订单编号）；
3.交易日期（注意交易日期在图中可能被标记为“结账日期”、“结账时间”、“下单时间”、“消费时间”或“日期”字段，请识别交易日期中的年月日，如果没识别到年份则默认为2024年，输出格式为 “年-月-日”）；
4.交易金额（注意交易金额在图中可能被标记为“可开发票金额”、“实付金额”、“实收金额”、“账单实收”、“收款金额”、“支付合计”或“消费合计”字段，请根据此字段顺序提取一个交易金额）。
注意将你的响应格式化为 JSON 对象，以“商户名称”、“订单编号”、“交易日期”、“交易金额” 4 个字段作为键，每个键对应一个值，如果信息不存在，请使用 “-” 作为值。
"""








print(prompt_txt)


def model_result(img_path):
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    query = prompt_txt
    image = Image.open(img_path).convert('RGB')
    inputs = tokenizer.apply_chat_template([{"role": "user", "image": image, "content": query}],
                                           add_generation_prompt=True, tokenize=True, return_tensors="pt",
                                           return_dict=True)  # chat mode
    
    inputs = inputs.to(device)

    gen_kwargs = {"max_length": 2500, "do_sample": True, "top_k": 1}
    with torch.no_grad():
        outputs = model.generate(**inputs, **gen_kwargs)
        outputs = outputs[:, inputs['input_ids'].shape[1]:]
        result_=tokenizer.decode(outputs[0])
    return result_
"""
img_path="a1.jpg"
result_=model_result(img_path)
print(result_)

files = os.listdir("test_img")
print(files)
for file in files:
    file_path = os.path.join("/rainbow/other/glm_4v_9b/test_img", file)
    print("file_path:",file_path)
    result_=model_result(file_path)
    print(result_)
"""